#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2021 wanghch <wanghch@wanghch-pc>
#
# Distributed under terms of the MIT license.

"""

"""
from feature_extract import *
import tensorflow as tf
from tf_utils import gen_tf_example
import os

today = datetime.datetime.today()
parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('-s', '--strategy', type=str, default="bottom_top", help='strategy fun')

ARGS = parser.parse_args()

class IOWrapper(object):
    def __init__(self, out, ftype = 'csv'):
        self.ftype = ftype
        if ftype == "pb":
            self.f = tf.io.TFRecordWriter(out + "." + ftype)
        elif ftype == "csv":
            self.f = open(out + "." + ftype, "w")
    def write(self, features, target):
        if self.ftype == "pb":
            features['target'] = target
            x = gen_tf_example(features)
            self.f.write(x.SerializeToString())
        elif self.ftype == "csv":
            fstr = ",".join(list(map(str, features.tolist())))
            self.f.write(fstr + "," + str(target) +  "\n")


    def close(self):
        self.f.close()


def main():
    data_dir = "data/"
    export_dir = "export/"
    out_dir = "out/"
    datafiles = os.listdir(data_dir)
    ftype = "csv"

    for dfile in tqdm(datafiles):
        code, name = dfile.replace(".csv", "").split("_")
        abs_dfile = os.path.join(data_dir, dfile)
        df = pd.read_csv(abs_dfile)
        N = df.shape[0]
        df = add_extra_indicators(df)
        # df = add_industry(df, code)
        out = os.path.join(out_dir, dfile.replace(".csv", ""))
        fdata = []
        fw = None
        for index,row in df.iterrows():
            if index < 250 or index >= N - 10:
                continue

            strategy_func = eval(ARGS.strategy)
            target = strategy_func(df, index, row)

            if not target is None:
                if fw is None:
                    fw = IOWrapper(out, ftype)
                features = extract_feature(df, index)
                fw.write(features, target)
        if fw:
            fw.close()


if __name__ == '__main__':
    main()
